package learner.hmm;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Class that represent the alpha values cfr. the algorithm from slides. a_k(i)
 * is probability of an observation sequence and a current state s_i.
 * 
 * @author Geert
 * 
 */
public class Alpha {

	private Matrix<Integer, Integer> a;

	private Matrix<Integer, MovementVector> b;

	private Map<Integer, Double> pi;

	private Movement movement;

	// list of already calculated alpha's, the index represents the time
	// (k-variable on slide)
	// the map is from a state to a probability.
	private List<Map<Integer, Double>> alphaKs;

	public Alpha(Matrix<Integer, Integer> a,
			Matrix<Integer, MovementVector> b,
			Map<Integer, Double> pi, Movement movement) {
		this.a = a;
		this.b = b;
		this.pi = pi;
		this.movement = movement;
		alphaKs = new ArrayList<Map<Integer, Double>>();
		initializeAlpha();
	}

	private void initializeAlpha() {
		for (int i = 0; i < BaumWelchLearner.NB_OF_STATES; i++) {
			double alpha_1_i = pi.get(i) * b.get(i, movement.getMovementForTime(1));
			setAlphaValue(1, i, alpha_1_i);
		}
	}

	// TODO: slide fout? zie forward recursion for HMM (moet alpha_k+1(j) zijn?
	public double get(int k, int j) {
		// trivial case:
		if (alphaValueKnown(k, j))
			return getKnownAlphaValue(k, j);
		else { // recursion
			double sum = 0;
			for (int i = 0 ; i < BaumWelchLearner.NB_OF_STATES; i ++) {
				sum += get(k-1, i) * a.get(i,j); // a_j_i in formule (fout slide??)
			}
			setAlphaValue(k, j, sum * b.get(j, movement.getMovementForTime(k)));
			return sum * b.get(j, movement.getMovementForTime(k));
		}
	}

	private boolean alphaValueKnown(int k, int i) {
		if (alphaKs.size() < k)
			return false;
		if (!alphaKs.get(k -1).containsKey(i))
			return false;
		return true;
	}
	
	// k can range from 1 to nbof movements
	private void setAlphaValue(int k, int i, double value) {
		if (k > alphaKs.size() + 1)
			throw new IllegalStateException(
					"You can't know the value yet: known alphaKs= "
							+ alphaKs.size() + " given time= " + k);
		if (alphaValueKnown(k, i))
			throw new IllegalStateException("Alpha value already known");
		if (k == alphaKs.size() + 1) {
			Map<Integer, Double> alphas_k = new HashMap<Integer, Double>();
			alphas_k.put(i, value);
			alphaKs.add(alphas_k);
			
		} else {
			alphaKs.get(k - 1).put(i, value);
		}
	}
	
	private double getKnownAlphaValue(int k, int i) {
		return alphaKs.get(k-1).get(i);
	}

}
